1#!/usr/bin/env python3
2# yt-upscaler: keeps the Independent library's thumbnails sharp on a TV.
3# youtube only stores ~720p thumbnails, which jellyfin then stretches across
4# the whole screen as the item backdrop (soft + blocky). this re-fetches the
5# highest-res thumbnail from the image CDN (not bot-walled) and runs it through
6# Real-ESRGAN general-x4v3 (exported to ONNX, so no pickle is ever loaded) at
7# 4x on CPU, writing a crisp <video>.webp that jellyfin uses for card + backdrop.
8import io
9import json
10import os
11import re
12import sys
13import time
14import urllib.request
15
16import numpy as np
17import onnxruntime as ort
18from PIL import Image
19
20INDEP_DIR = os.environ.get("INDEP_DIR", "/media/jellyfin/Independent")
21STATE_DIR = os.environ.get("STATE_DIR", "/state")
22MODEL = os.environ.get("SR_MODEL", "/app/realesr-general-x4v3.onnx")
23MIN_H = int(os.environ.get("SR_MIN_H", "1400")) # >= this tall ⇒ already done
24CAP_H = int(os.environ.get("SR_CAP_H", "2160")) # cap output height (4k-ish)
25TILE = int(os.environ.get("SR_TILE", "256")) # tile size to bound memory
26THREADS = int(os.environ.get("SR_THREADS", "4")) # leave cpu for other services
27SCAN_INTERVAL = int(os.environ.get("SR_SCAN_INTERVAL", "1800"))
28VIDEO_EXTS = (".webm", ".mp4", ".mkv")
29UA = {"User-Agent": "Mozilla/5.0 (yt-upscaler; +https://paperclover.net)"}
30
31status = {"enabled": True, "running": False, "done": 0, "total": 0,
32 "current": "", "errors": 0}
33_sess = None
34
35
36def log(m):
37 print(m, flush=True)
38
39
40def write_status():
41 try:
42 tmp = os.path.join(STATE_DIR, "upscale-status.json.tmp")
43 with open(tmp, "w") as f:
44 json.dump(status, f)
45 os.replace(tmp, os.path.join(STATE_DIR, "upscale-status.json"))
46 except OSError:
47 pass
48
49
50def enabled():
51 try:
52 with open(os.path.join(STATE_DIR, "upscale-enabled.json")) as f:
53 return bool(json.load(f)["enabled"])
54 except FileNotFoundError:
55 return True
56
57
58def session():
59 global _sess
60 if _sess is None:
61 opts = ort.SessionOptions()
62 opts.intra_op_num_threads = THREADS
63 opts.inter_op_num_threads = 1
64 _sess = ort.InferenceSession(MODEL, sess_options=opts,
65 providers=["CPUExecutionProvider"])
66 return _sess
67
68
69def upscale(img, scale=4, pad=8):
70 # tiled 4x SR; overlap each tile by `pad` px and crop it back to hide seams
71 arr = np.asarray(img, dtype=np.float32) / 255.0
72 h, w, _ = arr.shape
73 out = np.zeros((h * scale, w * scale, 3), dtype=np.float32)
74 sess = session()
75 for y in range(0, h, TILE):
76 for x in range(0, w, TILE):
77 y0, x0 = max(0, y - pad), max(0, x - pad)
78 y1, x1 = min(h, y + TILE + pad), min(w, x + TILE + pad)
79 patch = arr[y0:y1, x0:x1].transpose(2, 0, 1)[None]
80 res = sess.run(None, {sess.get_inputs()[0].name: patch})[0][0].transpose(1, 2, 0)
81 th, tw = min(TILE, h - y) * scale, min(TILE, w - x) * scale
82 ty, tx = (y - y0) * scale, (x - x0) * scale
83 out[y * scale:y * scale + th, x * scale:x * scale + tw] = \
84 res[ty:ty + th, tx:tx + tw]
85 return Image.fromarray((out.clip(0, 1) * 255).round().astype("uint8"))
86
87
88def yt_id(base):
89 ij = base + ".info.json"
90 if os.path.exists(ij):
91 try:
92 with open(ij) as f:
93 vid = json.load(f).get("id")
94 if vid:
95 return vid
96 except (OSError, json.JSONDecodeError):
97 pass
98 try:
99 with open(base + ".nfo") as f:
100 m = re.search(r"<plot>(.*?)</plot>", f.read(), re.S)
101 if m:
102 u = re.search(r"(?:v=|youtu\.be/)([A-Za-z0-9_-]{11})", m.group(1))
103 return u.group(1) if u else None
104 except OSError:
105 pass
106 return None
107
108
109def fetch_cdn_webp(vid):
110 for q in ("maxresdefault", "sddefault", "hqdefault"):
111 try:
112 req = urllib.request.Request(
113 f"https://i.ytimg.com/vi_webp/{vid}/{q}.webp", headers=UA)
114 with urllib.request.urlopen(req, timeout=30) as r:
115 data = r.read()
116 if len(data) > 1000:
117 return Image.open(io.BytesIO(data)).convert("RGB")
118 except Exception:
119 continue
120 return None
121
122
123def img_height(path):
124 try:
125 with Image.open(path) as im:
126 return im.height
127 except Exception:
128 return 0
129
130
131def enhance(video_path):
132 base = os.path.splitext(video_path)[0]
133 webp, jpg = base + ".webp", base + ".jpg"
134 if os.path.exists(webp) and img_height(webp) >= MIN_H:
135 return "skip"
136 src = None
137 vid = yt_id(base)
138 if vid:
139 src = fetch_cdn_webp(vid)
140 if src is None:
141 for p in (webp, jpg):
142 if os.path.exists(p):
143 src = Image.open(p).convert("RGB")
144 break
145 if src is None:
146 return "no-source"
147 up = upscale(src)
148 if up.height > CAP_H:
149 up = up.resize((round(up.width * CAP_H / up.height), CAP_H), Image.LANCZOS)
150 tmp = webp + ".tmp"
151 up.save(tmp, "WEBP", quality=92, method=6)
152 os.replace(tmp, webp)
153 if os.path.exists(jpg):
154 os.remove(jpg)
155 return "done"
156
157
158def independent_videos():
159 out = []
160 for ch in sorted(os.listdir(INDEP_DIR)) if os.path.isdir(INDEP_DIR) else []:
161 cdir = os.path.join(INDEP_DIR, ch)
162 if not os.path.isdir(cdir):
163 continue
164 for f in sorted(os.listdir(cdir)):
165 if f.lower().endswith(VIDEO_EXTS):
166 out.append(os.path.join(cdir, f))
167 return out
168
169
170def main():
171 log(f"yt-upscaler started (model={MODEL}, threads={THREADS})")
172 while True:
173 status["enabled"] = enabled()
174 if not status["enabled"]:
175 status.update(running=False, current="")
176 write_status()
177 time.sleep(5)
178 continue
179 vids = independent_videos()
180 pending = [v for v in vids
181 if img_height(os.path.splitext(v)[0] + ".webp") < MIN_H]
182 status.update(total=len(vids), done=len(vids) - len(pending),
183 running=bool(pending), current="")
184 write_status()
185 if pending:
186 log(f"upscaling {len(pending)} thumbnail(s)…")
187 for v in pending:
188 if not enabled():
189 break
190 status["current"] = os.path.basename(v)
191 write_status()
192 t = time.time()
193 try:
194 r = enhance(v)
195 if r == "done":
196 log(f"upscaled ({time.time()-t:.0f}s): {os.path.basename(v)}")
197 elif r != "skip":
198 log(f"{r}: {os.path.basename(v)}")
199 except Exception as e:
200 status["errors"] += 1
201 log(f"failed: {os.path.basename(v)}: {e}")
202 status["done"] += 1
203 write_status()
204 time.sleep(1) # be polite to the rest of the box
205 status.update(running=False, current="")
206 write_status()
207 if not enabled():
208 continue
209 time.sleep(SCAN_INTERVAL)
210
211
212if __name__ == "__main__":
213 if sys.argv[1:] == ["--check"]:
214 assert upscale(Image.new("RGB", (8, 8), (30, 60, 90))).size == (32, 32)
215 else:
216 main()