Update app.py

app.py

@@ -174,8 +174,8 @@ def process_image_detection(image, target_label, surprise_rating):
        original_size = image.size
        print(f"Image size: {original_size}")
 
-       # Calculate relative font size based on image dimensions
-       base_fontsize = min(original_size) / 40
+       # Calculate relative font size
+       base_fontsize = min(original_size) / 80
 
        print("Loading models...")
        owlv2_processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
@@ -191,8 +191,12 @@ 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)
+       
        dpi = 300
-       figsize = (original_size[0] / dpi, original_size[1] / dpi)
        fig = plt.figure(figsize=figsize, dpi=dpi)
        ax = plt.Axes(fig, [0., 0., 1., 1.])
        fig.add_axes(ax)
@@ -246,22 +250,22 @@ def process_image_detection(image, target_label, surprise_rating):
                )
                ax.add_patch(rect)
 
+               # Add confidence score without background
                plt.text(
                    box[0], box[1] - base_fontsize,
                    f'{max_score:.2f}',
                    color='red',
                    fontsize=base_fontsize,
-                   fontweight='bold',
-                   bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=2)
+                   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',
-                   bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', pad=2),
                    verticalalignment='bottom'
                )
 
@@ -269,7 +273,6 @@ def process_image_detection(image, target_label, surprise_rating):
 
        print("Saving final image...")
        try:
-           # Save directly to buffer using savefig
            buf = io.BytesIO()
            fig.savefig(buf, 
                       format='png',
@@ -285,16 +288,17 @@ def process_image_detection(image, target_label, surprise_rating):
            if output_image.mode != 'RGB':
                output_image = output_image.convert('RGB')
            
-           # Resize to original size if needed
-           if output_image.size != original_size:
-               output_image = output_image.resize(original_size, Image.Resampling.LANCZOS)
+           # 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)
            
-           # Cleanup
            plt.close(fig)
            buf.close()
            
@@ -302,8 +306,6 @@ def process_image_detection(image, target_label, surprise_rating):
 
        except Exception as e:
            print(f"Save error details: {str(e)}")
-           print(f"Figure type: {type(fig)}")
-           print(f"Canvas type: {type(fig.canvas)}")
            raise
 
    except Exception as e: