joermd commited on
Commit
6f2b794
·
verified ·
1 Parent(s): bcb4492

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -4
app.py CHANGED
@@ -1,5 +1,4 @@
1
  import streamlit as st
2
- from loadimg import load_img
3
  from transformers import AutoModelForImageSegmentation
4
  import torch
5
  from torchvision import transforms
@@ -7,12 +6,16 @@ from PIL import Image
7
  import requests
8
  from io import BytesIO
9
 
10
- # إعداد الجهاز
 
11
  torch.set_float32_matmul_precision("high")
 
 
12
  birefnet = AutoModelForImageSegmentation.from_pretrained(
13
  "ZhengPeng7/BiRefNet", trust_remote_code=True
14
- )
15
 
 
16
  transform_image = transforms.Compose(
17
  [
18
  transforms.Resize((1024, 1024)),
@@ -24,7 +27,7 @@ transform_image = transforms.Compose(
24
  # دالة المعالجة
25
  def process(image):
26
  image_size = image.size
27
- input_images = transform_image(image).unsqueeze(0).to("cuda")
28
  with torch.no_grad():
29
  preds = birefnet(input_images)[-1].sigmoid().cpu()
30
  pred = preds[0].squeeze()
 
1
  import streamlit as st
 
2
  from transformers import AutoModelForImageSegmentation
3
  import torch
4
  from torchvision import transforms
 
6
  import requests
7
  from io import BytesIO
8
 
9
+ # إعداد الجهاز للعمل على CPU
10
+ device = torch.device("cpu")
11
  torch.set_float32_matmul_precision("high")
12
+
13
+ # تحميل النموذج
14
  birefnet = AutoModelForImageSegmentation.from_pretrained(
15
  "ZhengPeng7/BiRefNet", trust_remote_code=True
16
+ ).to(device)
17
 
18
+ # تحويل الصورة
19
  transform_image = transforms.Compose(
20
  [
21
  transforms.Resize((1024, 1024)),
 
27
  # دالة المعالجة
28
  def process(image):
29
  image_size = image.size
30
+ input_images = transform_image(image).unsqueeze(0).to(device)
31
  with torch.no_grad():
32
  preds = birefnet(input_images)[-1].sigmoid().cpu()
33
  pred = preds[0].squeeze()