MoritzMMuller commited on
Commit
1164227
·
verified ·
1 Parent(s): 33f8c2e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -10
app.py CHANGED
@@ -43,16 +43,21 @@ geolocator = Nominatim(user_agent="skin-dashboard", timeout = 10)
43
 
44
  @st.cache_resource
45
  def load_image_model(token: str):
46
- """
47
- Load the pre-trained skin cancer classification model using PyTorch.
48
- """
49
- try:
50
- extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
51
- model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
52
- return pipeline("image-classification", model=model, feature_extractor=extractor, framework="pt")
53
- except Exception as e:
54
- st.error(f"Error loading the model: {e}")
55
- return None
 
 
 
 
 
56
 
57
 
58
  @st.cache_resource
 
43
 
44
  @st.cache_resource
45
  def load_image_model(token: str):
46
+ return pipeline(
47
+ "image-classification",
48
+ feature_extractor=AutoFeatureExtractor.from_pretrained(
49
+ MODEL_NAME,
50
+ subfolder="Skin_Cancer-Image_Classification",
51
+ use_auth_token=token
52
+ ),
53
+ model=AutoModelForImageClassification.from_pretrained(
54
+ MODEL_NAME,
55
+ subfolder="Skin_Cancer-Image_Classification",
56
+ use_auth_token=token
57
+ ),
58
+ device=0 # or -1 for CPU
59
+ )
60
+
61
 
62
 
63
  @st.cache_resource