Update app.py

app.py

@@ -1,5 +1,3 @@
-from flask import Flask, render_template, request, redirect, url_for, jsonify, flash
-from werkzeug.utils import secure_filename
 from ultralytics import YOLO
 import cv2
 import os
@@ -7,21 +5,20 @@ import numpy as np
 import uuid
 import time
 import logging
-
-app = Flask(__name__)
-app.secret_key = os.urandom(24)
-app.config['UPLOAD_FOLDER'] = 'static/uploads'
-app.config['RESULT_FOLDER'] = 'static/results'
-app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'webp', 'gif'}
-app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16 MB max upload
+import gradio as gr
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Create folders if they don't exist
-os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
-os.makedirs(app.config['RESULT_FOLDER'], exist_ok=True)
+UPLOAD_FOLDER = 'static/uploads'
+RESULT_FOLDER = 'static/results'
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp', 'gif'}
+MAX_CONTENT_LENGTH = 16 * 1024 * 1024  # 16 MB max upload
+
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(RESULT_FOLDER, exist_ok=True)
 
 # Load the model
 try:
@@ -31,20 +28,17 @@ except Exception as e:
     logger.error(f"Error loading model: {str(e)}")
     model = None
 
-def allowed_file(filename):
-    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']
-
-def count_fruits(image_path, model, conf=0.25):
+def count_fruits(image, conf=0.25):
     """
-    Runs inference on the image, counts fruits per class, and returns both the counts and results.
+    Runs inference on the image, counts fruits per class, and returns both the counts and visualized results.
     """
     try:
         # Run inference
-        results = model.predict(image_path, conf=conf)[0]
-
+        results = model.predict(image, conf=conf)[0]
+        
         # Initialize fruit counter dictionary
         fruit_counts = {}
-
+        
         # Iterate over each detected box
         for box in results.boxes:
             class_id = int(box.cls)
@@ -60,102 +54,81 @@ def count_fruits(image_path, model, conf=0.25):
                 class_with_confidence = f"{class_name}"
                 
             fruit_counts[class_with_confidence] = fruit_counts.get(class_with_confidence, 0) + 1
-
-        return fruit_counts, results, None
-    except Exception as e:
-        logger.error(f"Error in fruit detection: {str(e)}")
-        return {}, None, str(e)
-
-@app.route('/')
-def index():
-    return render_template('index.html')
-
-@app.route('/upload', methods=['POST'])
-def upload_file():
-    if 'file' not in request.files:
-        return jsonify({'error': 'No file part'}), 400
-    
-    file = request.files['file']
-    
-    if file.filename == '':
-        return jsonify({'error': 'No file selected'}), 400
-    
-    if not allowed_file(file.filename):
-        return jsonify({'error': 'File type not allowed'}), 400
-    
-    try:
-        # Generate a unique filename
-        unique_filename = f"{uuid.uuid4()}_{secure_filename(file.filename)}"
-        file_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
-        file.save(file_path)
         
-        # Return the filename so the frontend can use it for the next request
-        return jsonify({
-            'status': 'success',
-            'filename': unique_filename,
-            'message': 'File uploaded successfully'
-        })
+        # Generate result image with detections
+        results_image = results.plot()
+        
+        # Create count summary
+        count_summary = "\n".join([f"{fruit}: {count}" for fruit, count in fruit_counts.items()])
         
+        return results_image, count_summary
     except Exception as e:
-        logger.error(f"Error uploading file: {str(e)}")
-        return jsonify({'error': str(e)}), 500
+        logger.error(f"Error in fruit detection: {str(e)}")
+        return None, f"Error processing image: {str(e)}"
 
-@app.route('/process', methods=['POST'])
-def process_image():
-    data = request.json
-    filename = data.get('filename')
-    conf_threshold = float(data.get('confidence', 0.25))
-    
-    if not filename:
-        return jsonify({'error': 'No filename provided'}), 400
-    
-    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-    
-    if not os.path.exists(file_path):
-        return jsonify({'error': 'File not found'}), 404
+def process_image(image, confidence):
+    """
+    Process the uploaded image and return results for Gradio interface
+    """
+    if image is None:
+        return None, "No image uploaded"
     
     try:
-        # Process the image
+        # Process the image with the model
         start_time = time.time()
-        counts, results, error = count_fruits(file_path, model, conf=conf_threshold)
+        result_image, count_summary = count_fruits(image, conf=confidence)
+        processing_time = time.time() - start_time
         
-        if error:
-            return jsonify({'error': error}), 500
+        # Add processing time to summary
+        full_summary = f"{count_summary}\n\nProcessing time: {processing_time:.2f} seconds"
         
-        # Generate and save result image
-        results_image = results.plot()  # This returns a numpy array with the detections drawn
-        result_filename = f"result_{filename}"
-        result_path = os.path.join(app.config['RESULT_FOLDER'], result_filename)
+        return result_image, full_summary
+    except Exception as e:
+        logger.error(f"Error processing image: {str(e)}")
+        return None, f"Error: {str(e)}"
+
+# Create Gradio interface
+def create_interface():
+    with gr.Blocks(title="Fruit Detection") as demo:
+        gr.Markdown("# 🍎 Fruit Detection and Counting 🍌")
+        gr.Markdown("Upload an image to detect and count fruits")
         
-        # Save the result image
-        cv2.imwrite(result_path, results_image)
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(type="numpy", label="Input Image")
+                confidence = gr.Slider(
+                    minimum=0.1, 
+                    maximum=1.0, 
+                    value=0.25, 
+                    step=0.05, 
+                    label="Confidence Threshold"
+                )
+                submit_btn = gr.Button("Detect Fruits")
+            
+            with gr.Column():
+                output_image = gr.Image(type="numpy", label="Detection Results")
+                output_text = gr.Textbox(label="Fruit Counts")
         
-        # Get processing stats from results
-        processing_time = results.speed
-        total_time = time.time() - start_time
+        submit_btn.click(
+            fn=process_image,
+            inputs=[input_image, confidence],
+            outputs=[output_image, output_text]
+        )
         
-        return jsonify({
-            'status': 'success',
-            'counts': counts,
-            'original_image': filename,
-            'result_image': result_filename,
-            'processing_time': {
-                'preprocess': processing_time.get('preprocess', 'N/A'),
-                'inference': processing_time.get('inference', 'N/A'),
-                'postprocess': processing_time.get('postprocess', 'N/A'),
-                'total': round(total_time * 1000, 2)
-            }
-        })
+        gr.Markdown("## How to use")
+        gr.Markdown("1. Upload an image containing fruits")
+        gr.Markdown("2. Adjust the confidence threshold if needed")
+        gr.Markdown("3. Click 'Detect Fruits' to process the image")
+        gr.Markdown("4. View the detected fruits and counts")
         
-    except Exception as e:
-        logger.error(f"Error processing image: {str(e)}")
-        return jsonify({'error': str(e)}), 500
+        if model is None:
+            gr.Warning("⚠️ Model failed to load. The application might not work correctly.")
+    
+    return demo
 
-# Health check endpoint
-@app.route('/health')
-def health_check():
-    return jsonify({'status': 'ok', 'model_loaded': model is not None})
+# Launch the Gradio app
+demo = create_interface()
 
-if __name__ == '__main__':
-    port = int(os.environ.get('PORT', 10000))
-    app.run(host='0.0.0.0', port=port, debug=True)
+if __name__ == "__main__":
+    # Gradio will use the default port for HF Spaces
+    demo.launch()