Update app.py

app.py

@@ -5,23 +5,23 @@ import requests
 from io import BytesIO
 import threading
 import time
-import json
 
-# Load models and processor only once using session state
+# Load models and processor only once using Streamlit session state
 if 'models_loaded' not in st.session_state:
-    # Image processor
     st.session_state.image_processor = AutoImageProcessor.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-topwear")
-    # Topwear model
     st.session_state.top_wear_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-topwear")
-    # Pattern model
     st.session_state.pattern_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-pattern-rgb")
-    # Print model
     st.session_state.print_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-print")
-    # Sleeve length model
     st.session_state.sleeve_length_model = AutoModelForImageClassification.from_pretrained("vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-sleeve-length")
     st.session_state.models_loaded = True
 
-# Functions for predictions
+# Define image processing and classification functions
+def imageprocessing(url):
+    response = requests.get(url)
+    image = Image.open(BytesIO(response.content))
+    encoding = st.session_state.image_processor(image.convert("RGB"), return_tensors="pt")
+    return encoding, image
+
 def topwear(encoding):
     outputs = st.session_state.top_wear_model(**encoding)
     predicted_class_idx = outputs.logits.argmax(-1).item()
@@ -42,15 +42,10 @@ def sleevelengths(encoding):
     predicted_class_idx = outputs.logits.argmax(-1).item()
     return st.session_state.sleeve_length_model.config.id2label[predicted_class_idx]
 
-def imageprocessing(url):
-    response = requests.get(url)
-    image = Image.open(BytesIO(response.content))
-    encoding = st.session_state.image_processor(image.convert("RGB"), return_tensors="pt")
-    return encoding, image
+# Run all models in parallel
+def pipes(image_url):
+    encoding, image = imageprocessing(image_url)
 
-def pipes(imagepath):
-    encoding, image = imageprocessing(imagepath)
-    # Using threading for faster results
     results = [None] * 4
 
     def update_results(index, func):
@@ -68,21 +63,27 @@ def pipes(imagepath):
     for thread in threads:
         thread.join()
 
-    dicts = {"top": results[0], "pattern": results[1], "print": results[2], "sleeve_length": results[3]}
-    return dicts, image
+    result_dict = {
+        "topwear": results[0],
+        "pattern": results[1],
+        "print": results[2],
+        "sleeve_length": results[3]
+    }
+
+    return result_dict, image
 
 # Streamlit app UI
 st.title("Clothing Classification Pipeline")
 
 image_url = st.text_input("Enter Image URL")
-
 if image_url:
     start_time = time.time()
-    results, img = pipes(image_url)
-    st.image(img.resize((200, 200)), caption="Uploaded Image", use_column_width=False)
-    
-    # Display results as JSON
-    st.write("Classification Results (JSON):")
-    st.json(results)  # Output as JSON
     
-    st.write(f"Time taken: {time.time() - start_time:.2f} seconds")
+    try:
+        result, img = pipes(image_url)
+        st.image(img.resize((200, 200)), caption="Uploaded Image", use_column_width=False)
+        st.write("Classification Results (JSON):")
+        st.json(result)  # Display results in JSON format
+        st.write(f"Time taken: {time.time() - start_time:.2f} seconds")
+    except Exception as e:
+        st.error(f"Error processing the image: {str(e)}")