Update app.py

app.py

@@ -191,12 +191,11 @@ 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]
 
-       # Calculate aspect ratio and figure size
-       height, width = image.size[1], image.size[0]
-       aspect_ratio = width / height
-       figsize = (6 * aspect_ratio, 6)
-       
+       # Use original image dimensions for figure size
        dpi = 300
+       width, height = image.size
+       figsize = (width / dpi, height / dpi)
+       
        fig = plt.figure(figsize=figsize, dpi=dpi)
        ax = plt.Axes(fig, [0., 0., 1., 1.])
        fig.add_axes(ax)
@@ -250,7 +249,7 @@ def process_image_detection(image, target_label, surprise_rating):
                )
                ax.add_patch(rect)
 
-               # Add confidence score without background
+               # Only add the probability score
                plt.text(
                    box[0], box[1] - base_fontsize,
                    f'{max_score:.2f}',
@@ -259,26 +258,16 @@ def process_image_detection(image, target_label, surprise_rating):
                    fontweight='bold'
                )
 
-               # Add label and rating without background
-               plt.text(
-                   box[2] + base_fontsize / 2, box[1],
-                   f'Unexpected (Rating: {surprise_rating}/5)\n{target_label}',
-                   color='red',
-                   fontsize=base_fontsize,
-                   fontweight='bold',
-                   verticalalignment='bottom'
-               )
-
        plt.axis('off')
 
        print("Saving final image...")
        try:
            buf = io.BytesIO()
            fig.savefig(buf, 
-                      format='png',
-                      dpi=dpi,
-                      bbox_inches='tight',
-                      pad_inches=0)
+                     format='png',
+                     dpi=dpi,
+                     bbox_inches='tight',
+                     pad_inches=0)
            buf.seek(0)
            
            # Open as PIL Image
@@ -288,13 +277,7 @@ def process_image_detection(image, target_label, surprise_rating):
            if output_image.mode != 'RGB':
                output_image = output_image.convert('RGB')
            
-           # Calculate new size preserving aspect ratio
-           target_width = original_size[0]
-           target_height = int(target_width / aspect_ratio)
-           
-           # Resize maintaining aspect ratio
-           output_image = output_image.resize((target_width, target_height), Image.Resampling.LANCZOS)
-           
+           # Save to final buffer
            final_buf = io.BytesIO()
            output_image.save(final_buf, format='PNG', dpi=original_dpi)
            final_buf.seek(0)