Update app.py

app.py

@@ -3,8 +3,8 @@ from transformers import AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 import requests
 from io import BytesIO
-import threading
 import time
+import torch
 
 # Load models and processor only once using Streamlit session state
 if 'models_loaded' not in st.session_state:
@@ -17,74 +17,61 @@ if 'models_loaded' not in st.session_state:
 
 # Define image processing and classification functions
 def topwear(encoding):
-    outputs = st.session_state.top_wear_model(**encoding)
+    with torch.no_grad():
+        outputs = st.session_state.top_wear_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    st.write(st.session_state.top_wear_model.config.id2label[predicted_class_idx])
+    st.write(f"Top Wear: {st.session_state.top_wear_model.config.id2label[predicted_class_idx]}")
     return st.session_state.top_wear_model.config.id2label[predicted_class_idx]
 
 def patterns(encoding):
-    outputs = st.session_state.pattern_model(**encoding)
+    with torch.no_grad():
+        outputs = st.session_state.pattern_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    st.write(st.session_state.pattern_model.config.id2label[predicted_class_idx])
+    st.write(f"Pattern: {st.session_state.pattern_model.config.id2label[predicted_class_idx]}")
     return st.session_state.pattern_model.config.id2label[predicted_class_idx]
 
 def prints(encoding):
-    outputs = st.session_state.print_model(**encoding)
+    with torch.no_grad():
+        outputs = st.session_state.print_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    st.write(st.session_state.print_model.config.id2label[predicted_class_idx])
+    st.write(f"Print: {st.session_state.print_model.config.id2label[predicted_class_idx]}")
     return st.session_state.print_model.config.id2label[predicted_class_idx]
 
 def sleevelengths(encoding):
-    outputs = st.session_state.sleeve_length_model(**encoding)
+    with torch.no_grad():
+        outputs = st.session_state.sleeve_length_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    st.write(st.session_state.sleeve_length_model.config.id2label[predicted_class_idx])
+    st.write(f"Sleeve Length: {st.session_state.sleeve_length_model.config.id2label[predicted_class_idx]}")
     return st.session_state.sleeve_length_model.config.id2label[predicted_class_idx]
 
 def imageprocessing(image):
-    encoding = st.session_state.image_processor(image.convert("RGB"), return_tensors="pt")
+    encoding = st.session_state.image_processor(images=image, return_tensors="pt")
     return encoding
 
-# Define the function that will be used in each thread
-def call_model(func, encoding, results, index):
-    results[index] = func(encoding)
-
-# Run all models in parallel
+# Run all models sequentially
 def pipes(image):
     # Process the image once and reuse the encoding
     encoding = imageprocessing(image)
 
-    # Prepare a list to store the results from each thread
-    results = [None] * 4
-
-    # Create threads for each function call
-    threads = [
-        threading.Thread(target=call_model, args=(topwear, encoding, results, 0)),
-        threading.Thread(target=call_model, args=(patterns, encoding, results, 1)),
-        threading.Thread(target=call_model, args=(prints, encoding, results, 2)),
-        threading.Thread(target=call_model, args=(sleevelengths, encoding, results, 3)),
-    ]
-
-    # Start all threads
-    for thread in threads:
-        thread.start()
-
-    # Wait for all threads to finish
-    for thread in threads:
-        thread.join()
+    # Get results from each model
+    topwear_result = topwear(encoding)
+    pattern_result = patterns(encoding)
+    print_result = prints(encoding)
+    sleeve_length_result = sleevelengths(encoding)
 
     # Combine the results into a dictionary
-    dicts = {
-        "top": results[0],
-        "pattern": results[1],
-        "print": results[2],
-        "sleeve_length": results[3]
+    results = {
+        "top": topwear_result,
+        "pattern": pattern_result,
+        "print": print_result,
+        "sleeve_length": sleeve_length_result
     }
-    st.write(dicts)
-    return dicts
+    st.write(results)
+    return results
 
 # Streamlit app UI
 st.title("Clothing Classification Pipeline")