Backup

Signed-off-by: prescientmoon <git@moonythm.dev>

scripts/main.py

@@ -66,7 +66,7 @@ def import_charts_from_csv():
                         difficulty,
                         level,
                         int(note_count.replace(",", "").replace(".", "")),
-                        int(float(cc) * 100),
+                        int(round(float(cc) * 100)),
                         note_design if len(note_design) else None,
                     ),
                 )

src/commands/score.rs

@@ -1,9 +1,7 @@
 use std::fmt::Display;
 
 use crate::context::{Context, Error};
-use crate::score::{
-	jacket_rects, CreatePlay, ImageCropper, ImageDimensions, Play, RelativeRect, Score,
-};
+use crate::score::{CreatePlay, ImageCropper, Play, Score, ScoreKind};
 use crate::user::{discord_it_to_discord_user, User};
 use image::imageops::FilterType;
 use poise::serenity_prelude::{CreateAttachment, CreateEmbed, CreateMessage};
@@ -80,11 +78,8 @@ pub async fn magic(
 				let bytes = file.download().await?;
 				let format = image::guess_format(&bytes)?;
 
-				let image = image::load_from_memory_with_format(&bytes, format)?.resize(
-					1024,
-					1024,
-					FilterType::Nearest,
-				);
+				let image = image::load_from_memory_with_format(&bytes, format)?;
+				let mut image = image.resize(1024, 1024, FilterType::Nearest);
 				// }}}
 				// {{{ Detection
 				// Create cropper and run OCR
@@ -98,19 +93,44 @@ pub async fn magic(
 				// This makes OCR more likely to work
 				let mut ocr_image = image.grayscale().blur(1.);
 
+				// {{{ Kind
 				let edited = CreateReply::default()
 					.reply(true)
-					.content(format!("Image {}: reading difficulty", i + 1));
+					.content(format!("Image {}: reading kind", i + 1));
 				handle.edit(ctx, edited).await?;
 
-				let difficulty = match cropper.read_difficulty(&ocr_image) {
+				let kind = match cropper.read_score_kind(&ocr_image) {
 					// {{{ OCR error handling
 					Err(err) => {
 						error_with_image(
 							ctx,
 							&cropper.bytes,
 							&file.filename,
-							"Could not read score from picture",
+							"Could not read kind from picture",
+							&err,
+						)
+						.await?;
+
+						continue;
+					}
+					// }}}
+					Ok(k) => k,
+				};
+				// }}}
+				// {{{ Difficulty
+				let edited = CreateReply::default()
+					.reply(true)
+					.content(format!("Image {}: reading difficulty", i + 1));
+				handle.edit(ctx, edited).await?;
+
+				let difficulty = match cropper.read_difficulty(&ocr_image, kind) {
+					// {{{ OCR error handling
+					Err(err) => {
+						error_with_image(
+							ctx,
+							&cropper.bytes,
+							&file.filename,
+							"Could not read difficulty from picture",
 							&err,
 						)
 						.await?;
@@ -120,68 +140,37 @@ pub async fn magic(
 					// }}}
 					Ok(d) => d,
 				};
-
-				let song_by_jacket = cropper.read_jacket(ctx.data(), &image, difficulty).await;
+				// }}}
+				// {{{ Jacket & distribution
+				let mut jacket_rect = None;
+				let song_by_jacket = cropper
+					.read_jacket(ctx.data(), &mut image, kind, difficulty, &mut jacket_rect)
+					.await;
 				let note_distribution = cropper.read_distribution(&image)?;
-
+				// }}}
 				ocr_image.invert();
-
+				// {{{ Title
 				let edited = CreateReply::default()
 					.reply(true)
 					.content(format!("Image {}: reading title", i + 1));
 				handle.edit(ctx, edited).await?;
 
-				let song_by_name =
-					cropper.read_song(&ocr_image, &ctx.data().song_cache, difficulty);
-				let (song, chart) = match (song_by_jacket, song_by_name) {
-					// {{{ Both errors
-					(Err(err_jacket), Err(err_name)) => {
-						cropper.crop_image_to_bytes(
-							&image,
-							RelativeRect::from_aspect_ratio(
-								ImageDimensions::from_image(&image),
-								jacket_rects(),
-							)
-							.ok_or_else(|| "Could not find jacket area in picture")?
-							.to_absolute(),
-						)?;
-						error_with_image(
-							ctx,
-							&cropper.bytes,
-							&file.filename,
-							"Hey! I could not read the score in the provided picture.",
-							&format!(
-                                "This can mean one of three things:
-1. The image you provided is *not that of an Arcaea score
-2. The image you provided contains a newly added chart that is not in my database yet
-3. The image you provided contains character art that covers the chart name. When this happens, I try to make use of the jacket art in order to determine the chart. It is possible that I've never seen the jacket art for this particular song on this particular difficulty. Contact `@prescientmoon` on discord in order to resolve the issue for you & future users playing this chart!
+				let song_by_name = match kind {
+					ScoreKind::SongSelect => None,
+					ScoreKind::ScoreScreen => {
+						Some(cropper.read_song(&ocr_image, &ctx.data().song_cache, difficulty))
+					}
+				};
 
-Nerdy info:
-```
-Jacket error: {}
-Title error: {}
-```" ,
-								err_jacket, err_name
-							),
-						)
-						.await?;
-						continue;
-					}
-					// }}}
-					// {{{ Only jacket succeeded
-					(Ok(by_jacket), Err(err_name)) => {
-						println!("Could not read name with error: {}", err_name);
-						by_jacket
-					}
-					// }}}
+				let (song, chart) = match (song_by_jacket, song_by_name) {
 					// {{{ Only name succeeded
-					(Err(err_jacket), Ok(by_name)) => {
+					(Err(err_jacket), Some(Ok(by_name))) => {
 						println!("Could not recognise jacket with error: {}", err_jacket);
 						by_name
 					}
 					// }}}
 					// {{{ Both succeeded
-					(Ok(by_jacket), Ok(by_name)) => {
+					(Ok(by_jacket), Some(Ok(by_name))) => {
 						if by_name.0.id != by_jacket.0.id {
 							println!(
 								"Got diverging choices between '{}' and '{}'",
@@ -191,8 +180,58 @@ Title error: {}
 
 						by_jacket
 					} // }}}
+					// {{{ Only jacket succeeded
+					(Ok(by_jacket), err_name) => {
+						if let Some(err) = err_name {
+							println!("Could not read name with error: {:?}", err.unwrap_err());
+						}
+
+						by_jacket
+					}
+					// }}}
+					// {{{ Both errors
+					(Err(err_jacket), err_name) => {
+						if let Some(rect) = jacket_rect {
+							cropper.crop_image_to_bytes(&image, rect)?;
+							error_with_image(
+							ctx,
+							&cropper.bytes,
+							&file.filename,
+							"Hey! I could not read the score in the provided picture.",
+							&format!(
+                                "This can mean one of three things:
+1. The image you provided is *not that of an Arcaea score
+2. The image you provided contains a newly added chart that is not in my database yet
+3. The image you provided contains character art that covers the chart name. When this happens, I try to make use of the jacket art in order to determine the chart. Contact `@prescientmoon` on discord to try and resolve the issue!
+
+Nerdy info:
+```
+Jacket error: {}
+Title error: {:?}
+```" ,
+								err_jacket, err_name
+							),
+						)
+						.await?;
+						} else {
+							ctx.reply(format!(
+								"This is a weird error that should never happen...
+Nerdy info:
+```
+Jacket error: {}
+Title error: {:?}
+```",
+								err_jacket, err_name
+							))
+							.await?;
+						}
+						continue;
+					} // }}}
 				};
 
+				println!("{}", song.title);
+				// }}}
+				// {{{ Score
 				let edited = CreateReply::default()
 					.reply(true)
 					.content(format!("Image {}: reading score", i + 1));
@@ -216,7 +255,7 @@ Title error: {}
 						// }}}
 						Ok(scores) => scores,
 					};
-
+				// }}}
 				// {{{ Build play
 				let (score, maybe_fars, score_warning) = Score::resolve_ambiguities(
 					score_possibilities,

src/image.rs

@@ -0,0 +1,53 @@
+use image::{DynamicImage, GenericImage, GenericImageView};
+
+use crate::bitmap::{Position, Rect};
+
+fn unsigned_in_bounds(image: &DynamicImage, x: i32, y: i32) -> bool {
+	x >= 0 && y >= 0 && image.in_bounds(x as u32, y as u32)
+}
+
+/// Performs a horizontal shear operation, without performing anti-aliasing
+pub fn xshear(image: &mut DynamicImage, rect: Rect, center: Position, shear: f32) {
+	let width = rect.width as i32;
+	for y in rect.y..rect.y + rect.height as i32 {
+		let skew = (shear * ((y - center.1) as f32)) as i32;
+		for i in rect.x..rect.x + width {
+			let x = if skew < 0 { i } else { rect.x + width - 1 - i };
+
+			if unsigned_in_bounds(image, x, y) {
+				let pixel = image.get_pixel(x as u32, y as u32);
+				if unsigned_in_bounds(image, x + skew, y) {
+					image.put_pixel((x + skew) as u32, y as u32, pixel);
+				};
+			};
+		}
+	}
+}
+
+/// Performs a horizontal shear operation, without performing anti-aliasing
+pub fn yshear(image: &mut DynamicImage, rect: Rect, center: Position, shear: f32) {
+	let height = rect.height as i32;
+	for x in rect.x..rect.x + rect.height as i32 {
+		let skew = (shear * ((x - center.0) as f32)) as i32;
+		for i in rect.y..rect.y + height {
+			let y = if skew < 0 { i } else { rect.y + height - 1 - i };
+
+			if unsigned_in_bounds(image, x, y) {
+				let pixel = image.get_pixel(x as u32, y as u32);
+				if unsigned_in_bounds(image, x, y + skew) {
+					image.put_pixel(x as u32, (y + skew) as u32, pixel);
+				};
+			};
+		}
+	}
+}
+
+/// Performs a rotation as a series of three shear operations
+/// Does not perform anti-aliasing.
+pub fn rotate(image: &mut DynamicImage, rect: Rect, center: Position, angle: f32) {
+	let alpha = -f32::tan(angle);
+	let beta = f32::sin(angle);
+	xshear(image, rect, center, alpha);
+	yshear(image, rect, center, beta);
+	xshear(image, rect, center, alpha);
+}

src/main.rs

@@ -9,6 +9,7 @@ mod bitmap;
 mod chart;
 mod commands;
 mod context;
+mod image;
 mod jacket;
 mod score;
 mod user;

src/score.rs

@@ -7,14 +7,17 @@ use std::sync::OnceLock;
 
 use edit_distance::edit_distance;
 use image::{imageops::FilterType, DynamicImage, GenericImageView};
+use num::integer::Roots;
 use num::{traits::Euclid, Rational64};
 use poise::serenity_prelude::{
 	Attachment, AttachmentId, CreateAttachment, CreateEmbed, CreateEmbedAuthor, Timestamp,
 };
 use tesseract::{PageSegMode, Tesseract};
 
+use crate::bitmap::Rect;
 use crate::chart::{Chart, Difficulty, Song, SongCache};
 use crate::context::{Error, UserContext};
+use crate::image::rotate;
 use crate::jacket::IMAGE_VEC_DIM;
 use crate::user::User;
 
@@ -566,6 +569,12 @@ impl Play {
 			None => None,
 		};
 
+		println!("Rating {:?}", self.score.play_rating(chart.chart_constant));
+		println!(
+			"Rating {:?}",
+			self.score.play_rating(chart.chart_constant) as f32 / 100.0
+		);
+
 		let mut embed = CreateEmbed::default()
 			.title(format!(
 				"{} [{:?} {}]",
@@ -576,7 +585,7 @@ impl Play {
 				"Rating",
 				format!(
 					"{:.2} (+?)",
-					(self.score.play_rating(chart.chart_constant)) as f32 / 100.
+					(self.score.play_rating(chart.chart_constant)) as f32 / 100.0
 				),
 				true,
 			)
@@ -646,6 +655,13 @@ mod score_tests {
 }
 // }}}
 // }}}
+// {{{ Score image kind
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ScoreKind {
+	SongSelect,
+	ScoreScreen,
+}
+// }}}
 // {{{ Image processing helpers
 // {{{ ImageDimensions
 #[derive(Debug, Clone, Copy)]
@@ -703,6 +719,11 @@ impl AbsoluteRect {
 			self.dimensions,
 		)
 	}
+
+	#[inline]
+	pub fn to_rect(&self) -> Rect {
+		Rect::new(self.x as i32, self.y as i32, self.width, self.height)
+	}
 }
 // }}}
 // {{{ RelativeRect
@@ -796,20 +817,32 @@ impl RelativeRect {
 	}
 }
 // }}}
+// }}}
 // {{{ Data points
+// {{{ Trait
+trait UIDataPoint {
+	fn aspect_ratio(&self) -> f32;
+}
+
+impl UIDataPoint for RelativeRect {
+	fn aspect_ratio(&self) -> f32 {
+		self.dimensions.aspect_ratio()
+	}
+}
+// }}}
 // {{{ Processing
-fn process_datapoints(rects: &mut Vec<RelativeRect>) {
-	rects.sort_by_key(|r| (r.dimensions.aspect_ratio() * 1000.0).floor() as u32);
+fn process_datapoints(points: &mut Vec<impl UIDataPoint>) {
+	points.sort_by_key(|r| (r.aspect_ratio() * 1000.0).floor() as u32);
 
 	// Filter datapoints that are close together
 	let mut i = 0;
-	while i < rects.len() - 1 {
-		let low = rects[i];
-		let high = rects[i + 1];
+	while i < points.len() - 1 {
+		let low = &points[i];
+		let high = &points[i + 1];
 
-		if (low.dimensions.aspect_ratio() - high.dimensions.aspect_ratio()).abs() < 0.001 {
+		if (low.aspect_ratio() - high.aspect_ratio()).abs() < 0.001 {
 			// TODO: we could interpolate here but oh well
-			rects.remove(i + 1);
+			points.remove(i + 1);
 		}
 
 		i += 1;
@@ -889,8 +922,8 @@ fn title_rects() -> &'static [RelativeRect] {
 	})
 }
 // }}}
-// {{{ Jacket
-pub fn jacket_rects() -> &'static [RelativeRect] {
+// {{{ Jacket (score screen)
+pub fn jacket_score_screen_rects() -> &'static [RelativeRect] {
 	static CELL: OnceLock<Vec<RelativeRect>> = OnceLock::new();
 	CELL.get_or_init(|| {
 		let mut rects: Vec<RelativeRect> = vec![
@@ -909,6 +942,19 @@ pub fn jacket_rects() -> &'static [RelativeRect] {
 	})
 }
 // }}}
+// {{{ Jacket (song select)
+pub fn jacket_song_select_rects() -> &'static [RelativeRect] {
+	static CELL: OnceLock<Vec<RelativeRect>> = OnceLock::new();
+	CELL.get_or_init(|| {
+		let mut rects: Vec<RelativeRect> = vec![
+			AbsoluteRect::new(465, 319, 730, 45, ImageDimensions::new(2532, 1170)).to_relative(),
+			AbsoluteRect::new(158, 411, 909, 74, ImageDimensions::new(2160, 1620)).to_relative(),
+		];
+		process_datapoints(&mut rects);
+		rects
+	})
+}
+// }}}
 // {{{ Note distribution
 pub fn note_distribution_rects() -> (
 	&'static [RelativeRect],
@@ -952,6 +998,18 @@ pub fn note_distribution_rects() -> (
 	})
 }
 // }}}
+// {{{ Score kind
+fn score_kind_rects() -> &'static [RelativeRect] {
+	static CELL: OnceLock<Vec<RelativeRect>> = OnceLock::new();
+	CELL.get_or_init(|| {
+		let mut rects: Vec<RelativeRect> = vec![
+			AbsoluteRect::new(237, 16, 273, 60, ImageDimensions::new(2532, 1170)).to_relative(),
+			AbsoluteRect::new(19, 15, 273, 60, ImageDimensions::new(2160, 1620)).to_relative(),
+		];
+		process_datapoints(&mut rects);
+		rects
+	})
+}
 // }}}
 // }}}
 // {{{ Recognise chart
@@ -1084,13 +1142,9 @@ pub struct ImageCropper {
 }
 
 impl ImageCropper {
-	pub fn crop_image_to_bytes(
-		&mut self,
-		image: &DynamicImage,
-		rect: AbsoluteRect,
-	) -> Result<(), Error> {
+	pub fn crop_image_to_bytes(&mut self, image: &DynamicImage, rect: Rect) -> Result<(), Error> {
 		self.bytes.clear();
-		let image = image.crop_imm(rect.x, rect.y, rect.width, rect.height);
+		let image = image.crop_imm(rect.x as u32, rect.y as u32, rect.width, rect.height);
 		let mut cursor = Cursor::new(&mut self.bytes);
 		image.write_to(&mut cursor, image::ImageFormat::Png)?;
 
@@ -1109,7 +1163,8 @@ impl ImageCropper {
 			&image.resize_exact(image.width(), image.height(), FilterType::Nearest),
 			RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), score_rects())
 				.ok_or_else(|| "Could not find score area in picture")?
-				.to_absolute(),
+				.to_absolute()
+				.to_rect(),
 		)?;
 
 		let mut results = vec![];
@@ -1230,12 +1285,21 @@ impl ImageCropper {
 	}
 	// }}}
 	// {{{ Read difficulty
-	pub fn read_difficulty(&mut self, image: &DynamicImage) -> Result<Difficulty, Error> {
+	pub fn read_difficulty(
+		&mut self,
+		image: &DynamicImage,
+		kind: ScoreKind,
+	) -> Result<Difficulty, Error> {
+		if kind == ScoreKind::SongSelect {
+			return Ok(Difficulty::FTR);
+		}
+
 		self.crop_image_to_bytes(
 			&image,
 			RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), difficulty_rects())
 				.ok_or_else(|| "Could not find difficulty area in picture")?
-				.to_absolute(),
+				.to_absolute()
+				.to_rect(),
 		)?;
 
 		let mut t = Tesseract::new(None, Some("eng"))?.set_image_from_mem(&self.bytes)?;
@@ -1263,6 +1327,40 @@ impl ImageCropper {
 		Ok(difficulty)
 	}
 	// }}}
+	// {{{ Read score kind
+	pub fn read_score_kind(&mut self, image: &DynamicImage) -> Result<ScoreKind, Error> {
+		self.crop_image_to_bytes(
+			&image,
+			RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), score_kind_rects())
+				.ok_or_else(|| "Could not find score kind area in picture")?
+				.to_absolute()
+				.to_rect(),
+		)?;
+
+		let mut t = Tesseract::new(None, Some("eng"))?.set_image_from_mem(&self.bytes)?;
+		t.set_page_seg_mode(PageSegMode::PsmRawLine);
+		t = t.recognize()?;
+
+		let text: &str = &t.get_text()?;
+		let text = text.trim().to_lowercase();
+
+		let conf = t.mean_text_conf();
+		if conf < 10 && conf != 0 {
+			Err(format!(
+				"Score kind text is not readable (confidence = {}, text = {}).",
+				conf, text
+			))?;
+		}
+
+		let result = if edit_distance(&text, "Result") < edit_distance(&text, "Select a song") {
+			ScoreKind::ScoreScreen
+		} else {
+			ScoreKind::SongSelect
+		};
+
+		Ok(result)
+	}
+	// }}}
 	// {{{ Read song
 	pub fn read_song<'a>(
 		&mut self,
@@ -1274,7 +1372,8 @@ impl ImageCropper {
 			&image,
 			RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), title_rects())
 				.ok_or_else(|| "Could not find title area in picture")?
-				.to_absolute(),
+				.to_absolute()
+				.to_rect(),
 		)?;
 
 		let mut t = Tesseract::new(None, Some("eng"))?
@@ -1304,15 +1403,45 @@ impl ImageCropper {
 	pub async fn read_jacket<'a>(
 		&mut self,
 		ctx: &'a UserContext,
-		image: &DynamicImage,
+		image: &mut DynamicImage,
+		kind: ScoreKind,
 		difficulty: Difficulty,
+		out_rect: &mut Option<Rect>,
 	) -> Result<(&'a Song, &'a Chart), Error> {
-		let rect =
-			RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), jacket_rects())
-				.ok_or_else(|| "Could not find jacket area in picture")?
-				.to_absolute();
+		let rect = RelativeRect::from_aspect_ratio(
+			ImageDimensions::from_image(image),
+			if kind == ScoreKind::ScoreScreen {
+				jacket_score_screen_rects()
+			} else {
+				jacket_song_select_rects()
+			},
+		)
+		.ok_or_else(|| "Could not find jacket area in picture")?
+		.to_absolute();
 
-		let cropped = image.view(rect.x, rect.y, rect.width, rect.height);
+		let cropped = if kind == ScoreKind::ScoreScreen {
+			*out_rect = Some(rect.to_rect());
+			image.view(rect.x, rect.y, rect.width, rect.height)
+		} else {
+			let angle = f32::atan2(rect.height as f32, rect.width as f32);
+			let side = rect.height + rect.width;
+			rotate(
+				image,
+				Rect::new(rect.x as i32, rect.y as i32, side, side),
+				(rect.x as i32, (rect.y + rect.height) as i32),
+				angle,
+			);
+
+			let len = (rect.width.pow(2) + rect.height.pow(2)).sqrt();
+
+			*out_rect = Some(Rect::new(
+				rect.x as i32,
+				(rect.y + rect.height) as i32,
+				len,
+				len,
+			));
+			image.view(rect.x, rect.y + rect.height, len, len)
+		};
 		let (distance, song_id) = ctx
 			.jacket_cache
 			.recognise(&*cropped)
@@ -1341,7 +1470,8 @@ impl ImageCropper {
 			&image,
 			RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), pure_rects)
 				.ok_or_else(|| "Could not find pure-rect area in picture")?
-				.to_absolute(),
+				.to_absolute()
+				.to_rect(),
 		)?;
 
 		t = t.set_image_from_mem(&self.bytes)?.recognize()?;
@@ -1352,7 +1482,8 @@ impl ImageCropper {
 			&image,
 			RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), far_rects)
 				.ok_or_else(|| "Could not find far-rect area in picture")?
-				.to_absolute(),
+				.to_absolute()
+				.to_rect(),
 		)?;
 
 		t = t.set_image_from_mem(&self.bytes)?.recognize()?;
@@ -1363,7 +1494,8 @@ impl ImageCropper {
 			&image,
 			RelativeRect::from_aspect_ratio(ImageDimensions::from_image(image), lost_rects)
 				.ok_or_else(|| "Could not find lost-rect area in picture")?
-				.to_absolute(),
+				.to_absolute()
+				.to_rect(),
 		)?;
 
 		t = t.set_image_from_mem(&self.bytes)?.recognize()?;