Merge pull request #59 from travisbrown/topic/twshoot-fix

Fixes for screenshots

src/bin/twshoot.rs

@@ -4,6 +4,8 @@ use image::DynamicImage;
 use std::path::PathBuf;
 use std::time::Duration;
 
+const LOADING_DELAY: Duration = Duration::from_millis(1500);
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let opts: Opts = Opts::parse();
@@ -41,7 +43,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             status_id,
             opts.width,
             opts.height,
-            Some(Duration::from_millis(500)),
+            Some(LOADING_DELAY),
         )
         .await?;
 

src/browser/twitter/mod.rs

@@ -67,6 +67,14 @@ pub async fn shoot_tweet_bytes(
         tokio::time::sleep(duration).await;
     }
 
+    // There may be a cookies layer. If so we hide it.
+    client
+        .execute(
+            "document.getElementById('layers').children[0].style.display = 'none';",
+            vec![],
+        )
+        .await?;
+
     client.screenshot().await
 }
 
@@ -92,6 +100,12 @@ pub async fn shoot_tweet(
 
 const RGBA_WHITE: Rgba<u8> = Rgba([255, 255, 255, 255]);
 
+// TODO: Figure out why this is necessary for finding the right edge in some cases.
+fn is_very_light_gray(pixel: &Rgba<u8>) -> bool {
+    let threshhold = 253;
+    pixel.0[0] >= threshhold && pixel.0[1] >= threshhold && pixel.0[2] >= threshhold
+}
+
 pub fn crop_tweet<I: GenericImageView<Pixel = Rgba<u8>>>(
     buffer: &I,
 ) -> Option<(u32, u32, u32, u32)> {
@@ -116,7 +130,9 @@ pub fn crop_tweet<I: GenericImageView<Pixel = Rgba<u8>>>(
 
     // Continue moving right to the second intersecting line.
     while i < w {
-        if buffer.get_pixel(i, 0) != RGBA_WHITE {
+        let pixel = buffer.get_pixel(i, 0);
+
+        if !is_very_light_gray(&pixel) {
             right_edge = Some(i - 1);
             break;
         }
@@ -127,22 +143,52 @@ pub fn crop_tweet<I: GenericImageView<Pixel = Rgba<u8>>>(
         .zip(right_edge)
         .zip(gray)
         .and_then(|((left, right), gray)| {
-            i = 0;
-
             let mut upper_edge = None;
             let mut lower_edge = None;
+            let mut i = 0;
+
+            // We no longer have a top border, so we find the top of the profile image and count up from there.
+            // This is a terrible hack and needs to be improved.
 
-            // Start at the top at the newly-discovered left edge, move down until the first line.
+            // Find the top of the text above the profile image.
             while i < h {
-                if buffer.get_pixel(left, i) != RGBA_WHITE {
-                    upper_edge = Some(i + 2);
-                    i += 2;
+                if (left..=right)
+                    .map(|j| buffer.get_pixel(j, i))
+                    .any(|p| p != RGBA_WHITE)
+                {
                     break;
                 }
                 i += 1;
             }
 
-            // And the next line, which represents the bottom of the tweet, including the actions.
+            // Find the base of the text.
+            while i < h {
+                if !(left..=right)
+                    .map(|j| buffer.get_pixel(j, i))
+                    .any(|p| p != RGBA_WHITE)
+                {
+                    break;
+                }
+                i += 1;
+            }
+            let text_base = i;
+
+            // Find the top of the profile image.
+            while i < h {
+                if (left..=right)
+                    .map(|j| buffer.get_pixel(j, i))
+                    .any(|p| p != RGBA_WHITE)
+                {
+                    break;
+                }
+                i += 1;
+            }
+
+            upper_edge = Some(text_base + (i - text_base) / 2);
+
+            let mut i = 0;
+
+            // The first line represents the bottom of the tweet, including the actions.
             while i < h {
                 if buffer.get_pixel(left, i) != RGBA_WHITE {
                     lower_edge = Some(i - 1);
@@ -152,14 +198,16 @@ pub fn crop_tweet<I: GenericImageView<Pixel = Rgba<u8>>>(
             }
 
             upper_edge.zip(lower_edge).and_then(|(upper, lower)| {
-                i = lower;
+                // We move up two pixels because of a new double line.
+                // This should be fairly robust, since the target will always be higher anyway.
+                i = lower - 2;
 
                 let middle = left + (right - left) / 2;
                 let mut base = None;
 
                 // Finally move up until you hit another gray line.
                 while i > 0 {
-                    if buffer.get_pixel(middle, i) == gray {
+                    if buffer.get_pixel(middle, i) != RGBA_WHITE {
                         base = Some(i - 2);
                         break;
                     }
@@ -187,11 +235,11 @@ mod tests {
         let examples = vec![
             (
                 "examples/images/703033780689199104-full.png",
-                Some((252, 108, 1195, 423)),
+                Some((253, 99, 1195, 494)),
             ),
             (
-                "examples/images/1302424011511529474-full.png",
-                Some((252, 108, 1195, 665)),
+                "examples/images/1503631923154984960-full.png",
+                Some((253, 99, 1195, 1184)),
             ),
         ];