2025-03-25 13:29:57 +08:00

99 lines
2.9 KiB
Python
Executable File

#!/usr/bin/env python3
import argparse
import os
import requests
import json
import subprocess
class RoiResearch(object):
def __init__(self, token=None):
self.token = token
def login(self, username, password):
url = "https://themblem.com/api/v1/login/"
data = {
"username": username,
"password": password,
}
response = requests.post(url, json=data)
token = response.json()['token']
self.token = token
return token
def fetch_scan_data(self, last_id=None):
x = "/api/v1/scan-data/?limit=1000"
while True:
url = "https://themblem.com" + x
print("fetching", url)
response = requests.get(url, headers={"Authorization": f"Token {self.token}"})
r = response.json()
meta = r['meta']
if meta['next'] is None:
break
x = meta['next']
for sd in r['objects']:
if last_id and sd['id'] <= last_id:
return
yield sd
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--token", '-t', default='3ebd8c33-f46e-4b06-bda8-4c0f5f5eb530', type=str)
parser.add_argument("--username", "-u", type=str)
parser.add_argument("--password", "-p", type=str)
return parser.parse_args()
def get_last_id(list_file):
ret = None
if not os.path.exists(list_file):
return ret
with open(list_file, "r") as f:
for line in f:
if not line.strip():
continue
id = json.loads(line)['id']
if not ret or id > ret:
ret = id
return ret
def get_all_ids(list_file):
ret = []
with open(list_file, "r") as f:
for line in f:
ret.append(json.loads(line)['id'])
return ret
def main():
args = parse_args()
rr = RoiResearch(args.token)
if args.username and args.password:
print(rr.login(args.username, args.password))
list_file = "data/roi/list.txt"
last_id = get_last_id(list_file)
print("last_id", last_id)
with open(list_file, "a") as f:
for sd in rr.fetch_scan_data(last_id):
if not sd['image'] or not sd['labels']:
print("skipping", sd['id'])
continue
print("new id", sd['id'])
line = json.dumps({
'id': sd['id'],
'image': sd['image'],
'labels': sd['labels'],
})
f.write(line + "\n")
all_ids = get_all_ids(list_file)
for id in all_ids:
outd = f"data/roi/samples/{id}"
if os.path.exists(f"{outd}/{id}.json"):
print(f"skipping {id}, json already exists")
continue
cmd = f'../scripts/emcli get-scan-data {id} -o {outd}'
print(cmd)
subprocess.call(cmd, shell=True)
if __name__ == "__main__":
main()