Update infer.py
Browse files
infer.py
CHANGED
@@ -10,6 +10,21 @@ import ffmpeg
|
|
10 |
|
11 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
12 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
|
14 |
def infer_video(video_filepath: str, size_modifier: int) -> str:
|
15 |
model = RealESRGAN(device, scale=size_modifier)
|
|
|
10 |
|
11 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
12 |
|
13 |
+
def infer_image(img: Image.Image, size_modifier: int ) -> Image.Image:
|
14 |
+
if img is None:
|
15 |
+
raise Exception("Image not uploaded")
|
16 |
+
|
17 |
+
width, height = img.size
|
18 |
+
|
19 |
+
if width >= 5000 or height >= 5000:
|
20 |
+
raise Exception("The image is too large.")
|
21 |
+
|
22 |
+
model = RealESRGAN(device, scale=size_modifier)
|
23 |
+
model.load_weights(f'weights/RealESRGAN_x{size_modifier}.pth', download=False)
|
24 |
+
|
25 |
+
result = model.predict(img.convert('RGB'))
|
26 |
+
print(f"Image size ({device}): {size_modifier} ... OK")
|
27 |
+
return result
|
28 |
|
29 |
def infer_video(video_filepath: str, size_modifier: int) -> str:
|
30 |
model = RealESRGAN(device, scale=size_modifier)
|