Update app.py

app.py

@@ -17,37 +17,37 @@ if 'models_loaded' not in st.session_state:
     st.session_state.models_loaded = True
 
 # Define image processing and classification functions
-def topwear(encoding):
+def topwear(encoding, top_wear_model):
     with torch.no_grad():
-        outputs = st.session_state.top_wear_model(**encoding)
+        outputs = top_wear_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    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]
+    st.write(f"Top Wear: {top_wear_model.config.id2label[predicted_class_idx]}")
+    return top_wear_model.config.id2label[predicted_class_idx]
 
-def patterns(encoding):
+def patterns(encoding, pattern_model):
     with torch.no_grad():
-        outputs = st.session_state.pattern_model(**encoding)
+        outputs = pattern_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    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]
+    st.write(f"Pattern: {pattern_model.config.id2label[predicted_class_idx]}")
+    return pattern_model.config.id2label[predicted_class_idx]
 
-def prints(encoding):
+def prints(encoding, print_model):
     with torch.no_grad():
-        outputs = st.session_state.print_model(**encoding)
+        outputs = print_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    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]
+    st.write(f"Print: {print_model.config.id2label[predicted_class_idx]}")
+    return print_model.config.id2label[predicted_class_idx]
 
-def sleevelengths(encoding):
+def sleevelengths(encoding, sleeve_length_model):
     with torch.no_grad():
-        outputs = st.session_state.sleeve_length_model(**encoding)
+        outputs = sleeve_length_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    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]
+    st.write(f"Sleeve Length: {sleeve_length_model.config.id2label[predicted_class_idx]}")
+    return sleeve_length_model.config.id2label[predicted_class_idx]
 
 def imageprocessing(image):
     encoding = st.session_state.image_processor(images=image, return_tensors="pt")
@@ -58,13 +58,19 @@ def pipes(image):
     # Process the image once and reuse the encoding
     encoding = imageprocessing(image)
     
+    # Access models from session state before threading
+    top_wear_model = st.session_state.top_wear_model
+    pattern_model = st.session_state.pattern_model
+    print_model = st.session_state.print_model
+    sleeve_length_model = st.session_state.sleeve_length_model
+    
     # Define functions to run the models in parallel
     with concurrent.futures.ThreadPoolExecutor() as executor:
         futures = {
-            executor.submit(topwear, encoding): "topwear",
-            executor.submit(patterns, encoding): "patterns",
-            executor.submit(prints, encoding): "prints",
-            executor.submit(sleevelengths, encoding): "sleeve_length"
+            executor.submit(topwear, encoding, top_wear_model): "topwear",
+            executor.submit(patterns, encoding, pattern_model): "patterns",
+            executor.submit(prints, encoding, print_model): "prints",
+            executor.submit(sleevelengths, encoding, sleeve_length_model): "sleeve_length"
         }
 
         results = {}