Update app06.py

app06.py

@@ -1,17 +1,15 @@
 import streamlit as st
-from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor, SegformerConfig
+from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
 from PIL import Image
 import numpy as np
 import torch
-import os
 
-model_path = "/home/user/app/defectdetection/model"
-config = SegformerConfig.from_json_file(os.path.join(model_path, "config.json"))
+# Define the model path
+model_hub_path = "nazlicanto/model_defectdetection"
 
-model = SegformerForSemanticSegmentation(config=config)
-model.load_state_dict(torch.load(os.path.join(model_path, "pytorch_model.bin")))
-
-preprocessor = SegformerImageProcessor.from_pretrained(model_path, local_files_only=True)
+# Load the model and preprocessor
+model = SegformerForSemanticSegmentation.from_pretrained(model_hub_path)
+preprocessor = SegformerImageProcessor.from_pretrained(model_hub_path)
 
 st.title("PCB Defect Detection")