Update app.py

app.py

@@ -4,10 +4,11 @@ import numpy as np
 import cv2
 from PIL import Image
 import io
+import torch
 
 # Set page config
 st.set_page_config(
-    page_title="Stone Classification",
+    page_title="Stone Detection & Classification",
     page_icon="🪨",
     layout="wide"
 )
@@ -26,96 +27,107 @@ st.markdown("""
         text-align: center;
         padding: 2rem;
     }
-
     </style>
     """, unsafe_allow_html=True)
-    # .prediction-card {
-    #     padding: 2rem;
-    #     border-radius: 0.5rem;
-    #     background-color: #f0f2f6;
-    #     margin: 1rem 0;
-    # }
-    # .top-predictions {
-    #     margin-top: 2rem;
-    #     padding: 1rem;
-    #     background-color: white;
-    #     border-radius: 0.5rem;
-    #     box-shadow: 0 1px 3px rgba(0,0,0,0.12);
-    # }
-    # .prediction-bar {
-    #     display: flex;
-    #     align-items: center;
-    #     margin: 0.5rem 0;
-    # }
-    # .prediction-label {
-    #     width: 100px;
-    #     font-weight: 500;
-    # }
-@st.cache_resource
-def load_model():
-    """Load the trained model"""
-    return tf.keras.models.load_model('custom_model.h5')
 
-def preprocess_image(image):
-    """Preprocess the uploaded image"""
-    # # Convert to RGB if needed
-    # if image.mode != 'RGB':
-    #     image = image.convert('RGB')
-    
-    # Convert to numpy array
-    img_array = np.array(image)
+def resize_to_square(image):
+    """Resize image to square while maintaining aspect ratio"""
+    size = max(image.shape[0], image.shape[1])
+    new_img = np.zeros((size, size, 3), dtype=np.uint8)
     
-    # # Convert to RGB if needed
-    # if len(img_array.shape) == 2:  # Grayscale
-    #     img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
-    # elif img_array.shape[2] == 4:  # RGBA
-    #     img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
+    # Calculate position to paste original image
+    x_center = (size - image.shape[1]) // 2
+    y_center = (size - image.shape[0]) // 2
     
-    # # Preprocess image similar to training
-    # img_hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
-    # img_hsv[:, :, 2] = cv2.equalizeHist(img_hsv[:, :, 2])
-    # img_array = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)
+    # Copy the image into center of result image
+    new_img[y_center:y_center+image.shape[0], 
+            x_center:x_center+image.shape[1]] = image
     
-    # # Adjust brightness
-    # target_brightness = 150
-    # current_brightness = np.mean(img_array)
-    # alpha = target_brightness / (current_brightness + 1e-5)
-    # img_array = cv2.convertScaleAbs(img_array, alpha=alpha, beta=0)
+    return new_img
+
+@st.cache_resource
+def load_models():
+    """Load both object detection and classification models"""
+    # Load object detection model
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    object_detection_model = torch.load("fasterrcnn_resnet50_fpn_270824.pth", map_location=device)
+    object_detection_model.to(device)
+    object_detection_model.eval()
     
-    # # Apply Gaussian blur
-    # img_array = cv2.GaussianBlur(img_array, (5, 5), 0)
+    # Load classification model
+    classification_model = tf.keras.models.load_model('custom_model.h5')
     
-    # Resize
+    return object_detection_model, classification_model, device
+
+def perform_object_detection(image, model, device):
+    original_size = image.size
+    target_size = (256, 256)
+    frame_resized = cv2.resize(np.array(image), dsize=target_size, interpolation=cv2.INTER_AREA)
+    frame_rgb = cv2.cvtColor(frame_resized, cv2.COLOR_RGB2BGR).astype(np.float32)  
+    frame_rgb /= 255.0
+    frame_rgb = frame_rgb.transpose(2, 0, 1)
+    frame_rgb = torch.from_numpy(frame_rgb).float().unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = model(frame_rgb)
+
+    boxes = outputs[0]['boxes'].cpu().detach().numpy().astype(np.int32)
+    labels = outputs[0]['labels'].cpu().detach().numpy().astype(np.int32)
+    scores = outputs[0]['scores'].cpu().detach().numpy()
+
+    result_image = frame_resized.copy()
+    cropped_images = []
+    detected_boxes = []
+
+    for i in range(len(boxes)):
+        if scores[i] >= 0.75:
+            x1, y1, x2, y2 = boxes[i]
+            if (int(labels[i])-1) == 1 or (int(labels[i])-1) == 0:
+                color = (0, 0, 255)
+                label_text = 'Flame stone surface'
+
+                # Scale coordinates to original image size
+                original_h, original_w = original_size[::-1]
+                scale_h, scale_w = original_h / target_size[0], original_w / target_size[1]
+                x1_orig, y1_orig = int(x1 * scale_w), int(y1 * scale_h)
+                x2_orig, y2_orig = int(x2 * scale_w), int(y2 * scale_h)
+                
+                # Crop and process detected region
+                cropped_image = np.array(image)[y1_orig:y2_orig, x1_orig:x2_orig]
+                resized_crop = resize_to_square(cropped_image)
+                cropped_images.append(resized_crop)
+                detected_boxes.append((x1, y1, x2, y2))
+
+                # Draw bounding box
+                cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 3)
+                cv2.putText(result_image, label_text, (x1, y1 - 10), 
+                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+    return Image.fromarray(result_image), cropped_images, detected_boxes
+
+def preprocess_image(image):
+    """Preprocess the image for classification"""
+    img_array = np.array(image)
     img_array = cv2.resize(img_array, (256, 256))
-    
-    # Normalize
     img_array = img_array.astype('float32') / 255.0
-    
     return img_array
 
 def get_top_predictions(prediction, class_names, top_k=5):
     """Get top k predictions with their probabilities"""
-    # Get indices of top k predictions
     top_indices = prediction.argsort()[0][-top_k:][::-1]
-    
-    # Get corresponding class names and probabilities
     top_predictions = [
         (class_names[i], float(prediction[0][i]) * 100)
         for i in top_indices
     ]
-    
     return top_predictions
 
 def main():
-    # Title
-    st.title("🪨 Stone Classification")
-    st.write("Upload an image of a stone to classify its type")
+    st.title("🪨 Stone Detection & Classification")
+    st.write("Upload an image to detect and classify stone surfaces")
     
-    # Initialize session state for prediction if not exists
     if 'predictions' not in st.session_state:
         st.session_state.predictions = None
     
-    # Create two columns
     col1, col2 = st.columns(2)
     
     with col1:
@@ -123,53 +135,60 @@ def main():
         uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
         
         if uploaded_file is not None:
-            # Display uploaded image
             image = Image.open(uploaded_file)
             st.image(image, caption="Uploaded Image", use_column_width=True)
 
-            with st.spinner('Analyzing image...'):
+            with st.spinner('Processing image...'):
                 try:
-                    # Load model
-                    model = load_model()
+                    # Load both models
+                    object_detection_model, classification_model, device = load_models()
                     
-                    # Preprocess image
-                    processed_image = preprocess_image(image)
+                    # Perform object detection
+                    result_image, cropped_images, detected_boxes = perform_object_detection(
+                        image, object_detection_model, device
+                    )
                     
-                    # Make prediction
-                    prediction = model.predict(np.expand_dims(processed_image, axis=0))
+                    if not cropped_images:
+                        st.warning("No stone surfaces detected in the image")
+                        return
+                    
+                    # Display detection results
+                    st.subheader("Detection Results")
+                    st.image(result_image, caption="Detected Stone Surfaces", use_column_width=True)
+                    
+                    # Process each detected region
                     class_names = ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']
+                    all_predictions = []
                     
-                    # Get top 5 predictions
-                    top_predictions = get_top_predictions(prediction, class_names)
+                    for idx, cropped_image in enumerate(cropped_images):
+                        processed_image = preprocess_image(cropped_image)
+                        prediction = classification_model.predict(
+                            np.expand_dims(processed_image, axis=0)
+                        )
+                        top_predictions = get_top_predictions(prediction, class_names)
+                        all_predictions.append(top_predictions)
                     
                     # Store in session state
-                    st.session_state.predictions = top_predictions
+                    st.session_state.predictions = all_predictions
                     
                 except Exception as e:
-                    st.error(f"Error during prediction: {str(e)}")
+                    st.error(f"Error during processing: {str(e)}")
     
     with col2:
-        st.subheader("Prediction Results")
+        st.subheader("Classification Results")
         if st.session_state.predictions is not None:
-            # Create a card-like container for results
-            results_container = st.container()
-            with results_container:
-                # Display main prediction
-                st.markdown("<div class='prediction-card'>", unsafe_allow_html=True)
-                top_class, top_confidence = st.session_state.predictions[0]
-                st.markdown(f"### Primary Prediction: Grade {top_class}")
-                st.markdown(f"### Confidence: {top_confidence:.2f}%")
-                st.markdown("</div>", unsafe_allow_html=True)
+            for idx, predictions in enumerate(st.session_state.predictions):
+                st.markdown(f"### Region {idx + 1}")
                 
-                # Display confidence bar for top prediction
+                # Display main prediction
+                top_class, top_confidence = predictions[0]
+                st.markdown(f"**Primary Prediction: Grade {top_class}**")
+                st.markdown(f"**Confidence: {top_confidence:.2f}%**")
                 st.progress(top_confidence / 100)
                 
-                # Display top 5 predictions
-                st.markdown("### Top 5 Predictions")
-                st.markdown("<div class='top-predictions'>", unsafe_allow_html=True)
-                
-                # Create a Streamlit container for the predictions
-                for class_name, confidence in st.session_state.predictions:
+                # Display all predictions for this region
+                st.markdown("**Top 5 Predictions**")
+                for class_name, confidence in predictions:
                     col_label, col_bar, col_value = st.columns([2, 6, 2])
                     with col_label:
                         st.write(f"Grade {class_name}")
@@ -178,11 +197,9 @@ def main():
                     with col_value:
                         st.write(f"{confidence:.2f}%")
                 
-                st.markdown("</div>", unsafe_allow_html=True)
+                st.markdown("---")
         else:
-            st.info("Upload an image and click 'Predict' to see the results")
-    
-    # Footer
-    st.markdown("---")
+            st.info("Upload an image to see detection and classification results")
+
 if __name__ == "__main__":
     main()