Update app.py

app.py

@@ -86,6 +86,10 @@ def show_mask(mask, ax, random_color=False):
 def process_image_detection(image, target_label, surprise_rating):
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    # Get original image DPI and size
+    original_dpi = image.info.get('dpi', (72, 72))  # Default to 72 DPI if not specified
+    original_size = image.size
+
     owlv2_processor = Owlv2Processor.from_pretrained("google/owlv2-large-patch14")
     owlv2_model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-large-patch14").to(device)
 
@@ -101,9 +105,16 @@ def process_image_detection(image, target_label, surprise_rating):
     target_sizes = torch.tensor([image.size[::-1]]).to(device)
     results = owlv2_processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]
 
-    fig = plt.figure(figsize=(10, 10))
+    # Create figure with the exact pixel size of the original image
+    dpi = 100  # Base DPI for calculation
+    figsize = (original_size[0] / dpi, original_size[1] / dpi)
+    fig = plt.figure(figsize=figsize, dpi=dpi)
+
+    # Remove margins and spacing
+    ax = plt.Axes(fig, [0., 0., 1., 1.])
+    fig.add_axes(ax)
+
     plt.imshow(image)
-    ax = plt.gca()
 
     scores = results["scores"]
     if len(scores) > 0:
@@ -157,12 +168,26 @@ def process_image_detection(image, target_label, surprise_rating):
 
     plt.axis('off')
 
+    # Save with original resolution and DPI
     buf = io.BytesIO()
-    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+    plt.savefig(buf,
+                format='png',
+                dpi=dpi,
+                bbox_inches='tight',
+                pad_inches=0)
     buf.seek(0)
     plt.close()
 
-    return buf
+    # Open the buffer and create a new image with original properties
+    output_image = Image.open(buf)
+    output_image = output_image.resize(original_size, Image.Resampling.LANCZOS)
+
+    # Create a new buffer with the properly sized image
+    final_buf = io.BytesIO()
+    output_image.save(final_buf, format='PNG', dpi=original_dpi)
+    final_buf.seek(0)
+
+    return final_buf
 
 
 def process_and_analyze(image):
@@ -181,15 +206,14 @@ def process_and_analyze(image):
         gpt_response = analyze_image(image)
         response_data = json.loads(gpt_response)
 
-        analysis_text = f"Label: {response_data['label']}\nElement: {response_data['element']}\nRating: {response_data['rating']}/5"
-
         if response_data["label"].lower() == "surprising" and response_data["element"].lower() != "na":
             # Process image with detection models
             result_buf = process_image_detection(image, response_data["element"], response_data["rating"])
             result_image = Image.open(result_buf)
+            analysis_text = f"Label: {response_data['label']}\nElement: {response_data['element']}\nRating: {response_data['rating']}/5"
             return result_image, analysis_text
         else:
-            return image, f"{analysis_text}\nImage not surprising or no specific element found."
+            return image, "Not Surprising"
 
     except Exception as e:
         return None, f"Error processing image: {str(e)}"